<head>
  <style> body { margin: 0; } </style>

  <script type="importmap">{ "imports": {
    "react": "https://esm.sh/react",
    "react-dom": "https://esm.sh/react-dom/client"
  }}</script>

<!--  <script type="module">import * as React from 'react'; window.React = React;</script>-->
<!--  <script src="../../src/packages/react-force-graph-3d/dist/react-force-graph-3d.js" defer></script>-->
</head>

<body>
<div id="graph"></div>

<script src="//unpkg.com/@babel/standalone"></script>
<script type="text/jsx" data-type="module">
  import ForceGraph3D from 'https://esm.sh/react-force-graph-3d?external=react';
  import React, { useState, useMemo, useCallback } from 'react';
  import { createRoot } from 'react-dom';
  import { genRandomTree } from '../datasets/random-data.js';

  const ExpandableGraph = ({ graphData }) => {
    const rootId = 0;

    const nodesById = useMemo(() => {
      const nodesById = Object.fromEntries(graphData.nodes.map(node => [node.id, node]));

      // link parent/children
      graphData.nodes.forEach(node => {
        node.collapsed = node.id !== rootId;
        node.childLinks = [];
      });
      graphData.links.forEach(link => nodesById[link.source].childLinks.push(link));

      return nodesById;
    }, [graphData]);

    const getPrunedTree = useCallback(() => {
      const visibleNodes = [];
      const visibleLinks = [];
      (function traverseTree(node = nodesById[rootId]) {
        visibleNodes.push(node);
        if (node.collapsed) return;
        visibleLinks.push(...node.childLinks);
        node.childLinks
          .map(link => ((typeof link.target) === 'object') ? link.target : nodesById[link.target]) // get child node
          .forEach(traverseTree);
      })();

      return { nodes: visibleNodes, links: visibleLinks };
    }, [nodesById]);

    const [prunedTree, setPrunedTree] = useState(getPrunedTree());

    const handleNodeClick = useCallback(node => {
      node.collapsed = !node.collapsed; // toggle collapse state
      setPrunedTree(getPrunedTree())
    }, []);

    return <ForceGraph3D
      graphData={prunedTree}
      linkDirectionalParticles={2}
      nodeColor={node => !node.childLinks.length ? 'green' : node.collapsed ? 'red' : 'yellow'}
      onNodeClick={handleNodeClick}
    />;
  };

  createRoot(document.getElementById('graph')).render(
    <ExpandableGraph graphData={genRandomTree(600, true)}/>
  );
</script>
</body>